import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def run(rank_id, size):
    tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank_id
    print('before gather',' Rank ', rank_id, ' has data ', tensor)
    gather_list = [torch.zeros(2, dtype=torch.int64) for _ in range(4)]
    dist.all_gather(gather_list, tensor)
    # all_gather 是集体通信操作，作用是：

    # - 每个进程将自己的 tensor 发送给所有其他进程；
    # - 最终每个进程的 gather_list 会按进程 ID 顺序存储所有进程的 tensor 。
    # all_gather 的核心作用是让所有进程共享彼此的张量数据，最终每个进程都能获取到全局所有进程的数据，常用于需要全局信息同步的场景（如分布式数据拼接、全局统计等）。
    print('after gather',' Rank ', rank_id, ' has data ', tensor)
    print('after gather',' Rank ', rank_id, ' has gather list ', gather_list)


def init_process(rank_id, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank_id, world_size=size)
    fn(rank_id, size)


if __name__ == "__main__":
    size = 4
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
